#include "rsacrypto.h"
#include "base64.h"
#include <openssl/bio.h>
#include <openssl/rsa.h>
#include <openssl/pem.h>
#include <QDebug>

RsaCrypto::RsaCrypto(QObject *parent)
    : QObject{parent}
{}

// 构造对象并加载密钥文件(密钥以文件的形式存储在磁盘中)中的数据到内存中
RsaCrypto::RsaCrypto(QByteArray fileName,  KeyType type, QObject *parent) : QObject{parent}
{
    BIO* bio = BIO_new_file(fileName.data(), "rb");
    assert(bio != NULL);
    if (type == PublicKey) {
        PEM_read_bio_PUBKEY(bio, &m_pubkey, NULL, NULL);
    } else {
        PEM_read_bio_PrivateKey(bio, &m_prikey, NULL, NULL);
    }
    BIO_free(bio);
}

RsaCrypto::~RsaCrypto()
{
    if (m_pubkey) {
        EVP_PKEY_free(m_pubkey);
    }
    if (m_prikey) {
        EVP_PKEY_free(m_prikey);
    }
}

// 将密钥字符串解析为密钥类型
void RsaCrypto::parseStringToKey(QByteArray data, KeyType type)
{
    BIO* bio = BIO_new_mem_buf(data.data(), data.size());
    assert(bio != NULL);
    if (type == PublicKey) {
        PEM_read_bio_PUBKEY(bio, &m_pubkey, NULL, NULL);
    } else {
        PEM_read_bio_PrivateKey(bio, &m_prikey, NULL, NULL);
    }
    BIO_free(bio);
}

void RsaCrypto::generateRsaKey(KeyLength bits, QByteArray pub, QByteArray pri)
{
    // 1. 创建密钥上下文
    EVP_PKEY_CTX* ctx = EVP_PKEY_CTX_new_id(EVP_PKEY_RSA, NULL);
    // 2. 初始化ctx
    int ret = EVP_PKEY_keygen_init(ctx);
    assert(ret == 1);
    // 3. 指定密钥对的长度
    ret = EVP_PKEY_CTX_set_rsa_keygen_bits(ctx, bits);
    assert(ret == 1);
    // 4. 生成密钥对
    ret = EVP_PKEY_generate(ctx, &m_prikey);
    assert(ret == 1);
    // 5. 释放上下文
    EVP_PKEY_CTX_free(ctx);

    // 将密钥对从内存写入到磁盘文件
    // 写入私钥 （pem是一种数据格式）
    BIO* bio = BIO_new_file(pri.data(), "wb");
    ret = PEM_write_bio_PrivateKey(bio, m_prikey, NULL, NULL, 0, NULL, NULL);
    assert(ret == 1);
    BIO_flush(bio);
    BIO_free(bio);

    // 写入公钥
    bio = BIO_new_file(pub.data(), "wb");
    ret = PEM_write_bio_PUBKEY(bio, m_prikey);
    assert(ret == 1);
    BIO_flush(bio);
    BIO_free(bio);
}

// 通过公钥进行加密
QByteArray RsaCrypto::pubKeyEncrypt(QByteArray data)
{

    // 1. 创建密钥上下文
    EVP_PKEY_CTX* ctx = EVP_PKEY_CTX_new(m_pubkey, NULL);
    assert(ctx != NULL);

    // 2. 设置加密和填充模式
    int ret = EVP_PKEY_encrypt_init(ctx);
    assert(ret == 1);
    ret = EVP_PKEY_CTX_set_rsa_padding(ctx, RSA_PKCS1_OAEP_PADDING);
    assert(ret == 1);

    // 3. 使用公钥进行加密
    // 第一次调用函数是为了确定加密之后的密文长度
    size_t outLen = 0;
    ret = EVP_PKEY_encrypt(ctx, NULL, &outLen, reinterpret_cast<const unsigned char*>(data.data()), data.size());
    assert(ret == 1);

    QByteArray retStr;
    retStr.resize(outLen);

    ret = EVP_PKEY_encrypt(ctx, reinterpret_cast<unsigned char*>(retStr.data()), &outLen, reinterpret_cast<const unsigned char*>(data.data()), data.size());
    assert(ret == 1);
    assert(outLen <= retStr.size()); // 添加断言检查实际加密长度不超过预设长度
    // 4. 若实际加密后的长度小于预设长度，调整QByteArray大小
    if (outLen < retStr.size()) {
        retStr.resize(outLen);
    }


    // 进行Base64格式的转换
    Base64 base;
    retStr = base.encode(reinterpret_cast<char*>(retStr.data()), outLen);
    // retStr = retStr.toBase64();

    // 5. 释放资源
    EVP_PKEY_CTX_free(ctx);

    return retStr;
}

// 通过私钥进行解密
QByteArray RsaCrypto::privateKeyDecrypt(QByteArray data)
{
    // 1. 创建密钥上下文
    EVP_PKEY_CTX* ctx = EVP_PKEY_CTX_new(m_prikey, NULL);
    assert(ctx != NULL);

    // 2. 设置解密和填充模式
    int ret = EVP_PKEY_decrypt_init(ctx);
    assert(ret == 1);

    ret = EVP_PKEY_CTX_set_rsa_padding(ctx, RSA_PKCS1_OAEP_PADDING);
    assert(ret == 1);


    // 将进行Base64的编码的密文进行解码
    Base64 base;
    data = base.decode(data);
    // data = data.fromBase64(data);


    // 3. 使用私钥进行解密
    // 第一次调用函数是为了确定解密之后的明文长度
    size_t outLen = 0;
    ret = EVP_PKEY_decrypt(ctx, NULL, &outLen, reinterpret_cast<const unsigned char*>(data.data()), data.size());
    assert(ret == 1);
    QByteArray retStr;
    retStr.resize(outLen);

    ret = EVP_PKEY_decrypt(ctx, reinterpret_cast<unsigned char*>(retStr.data()), &outLen, reinterpret_cast<const unsigned char*>(data.data()), data.size());
    assert(ret == 1);
    assert(outLen <= retStr.size()); // 添加断言检查实际加密长度不超过预设长度
    // 4. 若实际加密后的长度小于预设长度，调整QByteArray大小
    if (outLen < retStr.size()) {
        retStr.resize(outLen);
    }

    // 5. 释放资源
    EVP_PKEY_CTX_free(ctx);

    return retStr;
}

// 数字签名
QByteArray RsaCrypto::sign(QByteArray data, QCryptographicHash::Algorithm hash)
{
    // 1. 计算哈希值
    QCryptographicHash h(hash);
    h.addData(data);
    QByteArray md = h.result();

    // 2. 创建上下文
    EVP_PKEY_CTX* ctx = EVP_PKEY_CTX_new(m_prikey, NULL);
    assert(ctx != NULL);

    // 3. 设置数字签名和填充模式
    int ret = EVP_PKEY_sign_init(ctx);
    assert(ret == 1);

    ret = EVP_PKEY_CTX_set_rsa_padding(ctx, RSA_PKCS1_PADDING);
    assert(ret == 1);

    // 4. 设置签名使用的哈希算法
    ret = EVP_PKEY_CTX_set_signature_md(ctx, m_hashMethods.value(hash)());
    assert(ret = 1);

    // 5. 使用私钥进行加密
    // 第一次调用函数是为了确定加密之后的密文长度
    size_t outLen = 0;
    ret = EVP_PKEY_sign(ctx, NULL, &outLen, reinterpret_cast<const unsigned char*>(md.data()), md.size());
    assert(ret == 1);
    QByteArray retStr;
    retStr.resize(outLen);
    ret = EVP_PKEY_sign(ctx, reinterpret_cast<unsigned char*>(retStr.data()), &outLen, reinterpret_cast<const unsigned char*>(md.data()), md.size());
    assert(ret == 1);

    // 将签名进行Base64编码
    Base64 base;
    retStr = base.encode(reinterpret_cast<char*>(retStr.data()), outLen);

    // retStr = retStr.toBase64();

    // 4. 释放资源
    EVP_PKEY_CTX_free(ctx);

    return retStr;
}

// 签名校验
bool RsaCrypto::verify(QByteArray sign, QByteArray data, QCryptographicHash::Algorithm hash)
{
    // 将进行Base64编码的签名解码
    Base64 base;
    sign = base.decode(sign);
    // sign = sign.fromBase64(sign);

    // 1. 计算哈希值
    QCryptographicHash h(hash);
    h.addData(data);
    QByteArray md = h.result();

    // 2. 创建上下文
    EVP_PKEY_CTX* ctx = EVP_PKEY_CTX_new(m_pubkey, NULL);
    assert(ctx != NULL);

    // 3. 设置签名校验和填充模式
    int ret = EVP_PKEY_verify_init(ctx);
    assert(ret == 1);

    ret = EVP_PKEY_CTX_set_rsa_padding(ctx, RSA_PKCS1_PADDING);
    assert(ret == 1);

    // 4. 设置校验签名的哈希算法
    ret = EVP_PKEY_CTX_set_signature_md(ctx, m_hashMethods.value(hash)());
    assert(ret = 1);

    // 5. 签名校验
    ret = EVP_PKEY_verify(ctx, reinterpret_cast<const unsigned char*>(sign.data()), sign.size()
                             , reinterpret_cast<const unsigned char*>(md.data()), md.size());

    // 6. 释放资源
    EVP_PKEY_CTX_free(ctx);

    if (ret == 1) {
        return true;
    }

    return false;
}
